import { Document } from 'mongoose';

export type TreeNodeResult<T> = T & { children: T[] };
export function arrToTree<T extends Document>(list: T[], field: keyof T): TreeNodeResult<T>[] {
    // 创建对象字典
    const map = list.reduce((pre, item) => {
        const tmpItem = item.toObject();
        pre.set(tmpItem._id?.toString(), tmpItem);
        return pre;
    }, new Map<string, T | TreeNodeResult<T>>());

    // 创建树
    return list.reduce(
        (pre, cur) => {
            const _id = cur._id?.toString();
            const parentId = cur[field]?.toString();
            if (parentId) {
                const parentItem = map.get(parentId) as any;
                if (!parentItem) return pre;

                parentItem.children ? parentItem.children.push(cur) : (parentItem.children = [cur]);
                return pre;
            }
            pre.push(map.get(_id) as any);

            return pre;
        },
        [] as (T & { children: T[] })[]
    );
}
